"""
字幕和配音时间轴管理模块
提供字幕生成、编辑和同步功能
"""

import re
import os
from typing import Dict, Any, Optional, List, Tuple
from dataclasses import dataclass
from enum import Enum
import json
import xml.etree.ElementTree as ET


class SubtitleFormat(Enum):
    """字幕格式"""
    SRT = "srt"
    VTT = "vtt"
    ASS = "ass"
    SBV = "sbv"
    XML = "xml"
    JSON = "json"


@dataclass
class SubtitleEntry:
    """字幕条目"""
    index: int
    start_time: float
    end_time: float
    text: str
    metadata: Dict[str, Any] = None


@dataclass
class Timeline:
    """时间轴"""
    entries: List[SubtitleEntry]
    duration: float = 0.0
    metadata: Dict[str, Any] = None


class SubtitleManager:
    """字幕管理器"""
    
    def __init__(self):
        self.format_parsers = {
            SubtitleFormat.SRT: self._parse_srt,
            SubtitleFormat.VTT: self._parse_vtt,
            SubtitleFormat.JSON: self._parse_json,
            SubtitleFormat.SBV: self._parse_sbv,
            SubtitleFormat.ASS: self._parse_ass
        }
        
        self.format_writers = {
            SubtitleFormat.SRT: self._write_srt,
            SubtitleFormat.VTT: self._write_vtt,
            SubtitleFormat.JSON: self._write_json,
            SubtitleFormat.SBV: self._write_sbv,
            SubtitleFormat.ASS: self._write_ass
        }
    
    def load_subtitles(
        self,
        file_path: str,
        format: Optional[SubtitleFormat] = None
    ) -> Timeline:
        """加载字幕文件"""
        try:
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"Subtitle file not found: {file_path}")
            
            # 自动检测格式
            if format is None:
                format = self._detect_format(file_path)
            
            # 解析字幕
            if format not in self.format_parsers:
                raise ValueError(f"Unsupported subtitle format: {format}")
            
            parser = self.format_parsers[format]
            entries = parser(file_path)
            
            # 计算总时长
            duration = max((entry.end_time for entry in entries), default=0.0)
            
            return Timeline(
                entries=entries,
                duration=duration,
                metadata={'source_file': file_path, 'format': format.value}
            )
            
        except Exception as e:
            raise RuntimeError(f"Failed to load subtitles: {str(e)}")
    
    def save_subtitles(
        self,
        timeline: Timeline,
        output_path: str,
        format: SubtitleFormat
    ) -> bool:
        """保存字幕文件"""
        try:
            if format not in self.format_writers:
                raise ValueError(f"Unsupported subtitle format: {format}")
            
            writer = self.format_writers[format]
            return writer(timeline, output_path)
            
        except Exception as e:
            raise RuntimeError(f"Failed to save subtitles: {str(e)}")
    
    def create_timeline_from_text(
        self,
        text: str,
        max_duration_per_entry: float = 5.0,
        words_per_second: float = 2.5
    ) -> Timeline:
        """从文本创建时间轴"""
        try:
            # 分段处理
            sentences = self._split_text(text)
            entries = []
            
            current_time = 0.0
            
            for i, sentence in enumerate(sentences):
                word_count = len(sentence.split())
                duration = min(word_count / words_per_second, max_duration_per_entry)
                
                entry = SubtitleEntry(
                    index=i + 1,
                    start_time=current_time,
                    end_time=current_time + duration,
                    text=sentence.strip()
                )
                
                entries.append(entry)
                current_time += duration
            
            return Timeline(
                entries=entries,
                duration=current_time,
                metadata={'auto_generated': True}
            )
            
        except Exception as e:
            raise RuntimeError(f"Failed to create timeline: {str(e)}")
    
    def sync_subtitles_to_audio(
        self,
        timeline: Timeline,
        audio_duration: float,
        adjustment_factor: float = 1.0
    ) -> Timeline:
        """同步字幕到音频时长"""
        try:
            entries = []
            
            for entry in timeline.entries:
                # 调整时间
                adjusted_start = entry.start_time * adjustment_factor
                adjusted_end = entry.end_time * adjustment_factor
                
                # 确保不超出音频时长
                if adjusted_start >= audio_duration:
                    continue
                
                if adjusted_end > audio_duration:
                    adjusted_end = audio_duration
                
                synced_entry = SubtitleEntry(
                    index=entry.index,
                    start_time=adjusted_start,
                    end_time=adjusted_end,
                    text=entry.text,
                    metadata=entry.metadata
                )
                
                entries.append(synced_entry)
            
            return Timeline(
                entries=entries,
                duration=audio_duration,
                metadata=timeline.metadata
            )
            
        except Exception as e:
            raise RuntimeError(f"Failed to sync subtitles: {str(e)}")
    
    def merge_timelines(
        self,
        timelines: List[Timeline],
        mode: str = "overlay"  # overlay, append, merge
    ) -> Timeline:
        """合并多个时间轴"""
        try:
            if not timelines:
                return Timeline(entries=[], duration=0.0)
            
            if mode == "append":
                return self._append_timelines(timelines)
            elif mode == "merge":
                return self._merge_timelines(timelines)
            else:  # overlay
                return self._overlay_timelines(timelines)
                
        except Exception as e:
            raise RuntimeError(f"Failed to merge timelines: {str(e)}")
    
    def _parse_srt(self, file_path: str) -> List[SubtitleEntry]:
        """解析SRT格式字幕"""
        entries = []
        
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # SRT格式：序号 -> 时间轴 -> 文本
        pattern = r'(\d+)\n(\d{2}:\d{2}:\d{2},\d{3}) --> (\d{2}:\d{2}:\d{2},\d{3})\n(.*?)(?=\n\d+\n|\Z)'
        matches = re.findall(pattern, content, re.DOTALL)
        
        for match in matches:
            index = int(match[0])
            start_time = self._parse_srt_time(match[1])
            end_time = self._parse_srt_time(match[2])
            text = match[3].strip()
            
            entry = SubtitleEntry(
                index=index,
                start_time=start_time,
                end_time=end_time,
                text=text
            )
            
            entries.append(entry)
        
        return entries
    
    def _parse_vtt(self, file_path: str) -> List[SubtitleEntry]:
        """解析VTT格式字幕"""
        entries = []
        
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # VTT格式：时间轴 -> 文本
        pattern = r'(\d{2}:\d{2}:\d{2}\.\d{3}) --> (\d{2}:\d{2}:\d{2}\.\d{3})\n(.*?)(?=\n\d{2}:\d{2}:\d{2}\.\d{3}|\Z)'
        matches = re.findall(pattern, content, re.DOTALL)
        
        for i, match in enumerate(matches):
            start_time = self._parse_vtt_time(match[0])
            end_time = self._parse_vtt_time(match[1])
            text = match[2].strip()
            
            entry = SubtitleEntry(
                index=i + 1,
                start_time=start_time,
                end_time=end_time,
                text=text
            )
            
            entries.append(entry)
        
        return entries
    
    def _parse_json(self, file_path: str) -> List[SubtitleEntry]:
        """解析JSON格式字幕"""
        entries = []
        
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        for item in data:
            entry = SubtitleEntry(
                index=item.get('index', 0),
                start_time=item.get('start_time', 0.0),
                end_time=item.get('end_time', 0.0),
                text=item.get('text', ''),
                metadata=item.get('metadata', {})
            )
            
            entries.append(entry)
        
        return entries
    
    def _parse_sbv(self, file_path: str) -> List[SubtitleEntry]:
        """解析SBV格式字幕"""
        entries = []
        
        with open(file_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
        
        i = 0
        while i < len(lines):
            if ',' in lines[i]:
                # 时间轴行
                time_line = lines[i].strip()
                text_line = lines[i + 1].strip() if i + 1 < len(lines) else ''
                
                start_str, end_str = time_line.split(',')
                start_time = self._parse_sbv_time(start_str)
                end_time = self._parse_sbv_time(end_str)
                
                entry = SubtitleEntry(
                    index=len(entries) + 1,
                    start_time=start_time,
                    end_time=end_time,
                    text=text_line
                )
                
                entries.append(entry)
                i += 2
            else:
                i += 1
        
        return entries
    
    def _parse_ass(self, file_path: str) -> List[SubtitleEntry]:
        """解析ASS格式字幕"""
        entries = []
        
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # 简化的ASS解析（实际ASS格式更复杂）
        pattern = r'Dialogue: .*?,(\d{1,2}):(\d{2}):(\d{2})\.(\d{2}),(\d{1,2}):(\d{2}):(\d{2})\.(\d{2}),.*?,.*?,.*?,.*?,.*?,.*?,.*?,(.*)'
        matches = re.findall(pattern, content)
        
        for i, match in enumerate(matches):
            start_time = self._parse_ass_time(match[:4])
            end_time = self._parse_ass_time(match[4:8])
            text = match[8].strip()
            
            entry = SubtitleEntry(
                index=i + 1,
                start_time=start_time,
                end_time=end_time,
                text=text
            )
            
            entries.append(entry)
        
        return entries
    
    def _write_srt(self, timeline: Timeline, output_path: str) -> bool:
        """写入SRT格式字幕"""
        try:
            with open(output_path, 'w', encoding='utf-8') as f:
                for entry in timeline.entries:
                    start_time = self._format_srt_time(entry.start_time)
                    end_time = self._format_srt_time(entry.end_time)
                    
                    f.write(f"{entry.index}\n")
                    f.write(f"{start_time} --> {end_time}\n")
                    f.write(f"{entry.text}\n\n")
            
            return True
            
        except Exception:
            return False
    
    def _write_vtt(self, timeline: Timeline, output_path: str) -> bool:
        """写入VTT格式字幕"""
        try:
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write("WEBVTT\n\n")
                
                for entry in timeline.entries:
                    start_time = self._format_vtt_time(entry.start_time)
                    end_time = self._format_vtt_time(entry.end_time)
                    
                    f.write(f"{start_time} --> {end_time}\n")
                    f.write(f"{entry.text}\n\n")
            
            return True
            
        except Exception:
            return False
    
    def _write_json(self, timeline: Timeline, output_path: str) -> bool:
        """写入JSON格式字幕"""
        try:
            data = []
            for entry in timeline.entries:
                item = {
                    'index': entry.index,
                    'start_time': entry.start_time,
                    'end_time': entry.end_time,
                    'text': entry.text
                }
                
                if entry.metadata:
                    item['metadata'] = entry.metadata
                
                data.append(item)
            
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
            
            return True
            
        except Exception:
            return False
    
    def _write_sbv(self, timeline: Timeline, output_path: str) -> bool:
        """写入SBV格式字幕"""
        try:
            with open(output_path, 'w', encoding='utf-8') as f:
                for entry in timeline.entries:
                    start_time = self._format_sbv_time(entry.start_time)
                    end_time = self._format_sbv_time(entry.end_time)
                    
                    f.write(f"{start_time},{end_time}\n")
                    f.write(f"{entry.text}\n\n")
            
            return True
            
        except Exception:
            return False
    
    def _write_ass(self, timeline: Timeline, output_path: str) -> bool:
        """写入ASS格式字幕"""
        try:
            with open(output_path, 'w', encoding='utf-8') as f:
                # ASS头部
                f.write("[Script Info]\n")
                f.write("Title: Subtitles\n")
                f.write("ScriptType: v4.00+\n\n")
                
                f.write("[V4+ Styles]\n")
                f.write("Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding\n")
                f.write("Style: Default,Arial,20,&H00FFFFFF,&H000000FF,&H00000000,&H00000000,0,0,0,0,100,100,0,0,1,0,0,2,0,0,0,1\n\n")
                
                f.write("[Events]\n")
                f.write("Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text\n")
                
                for entry in timeline.entries:
                    start_time = self._format_ass_time(entry.start_time)
                    end_time = self._format_ass_time(entry.end_time)
                    
                    f.write(f"Dialogue: 0,{start_time},{end_time},Default,,0,0,0,,{entry.text}\n")
            
            return True
            
        except Exception:
            return False
    
    def _detect_format(self, file_path: str) -> SubtitleFormat:
        """自动检测字幕格式"""
        ext = os.path.splitext(file_path)[1].lower()
        
        format_map = {
            '.srt': SubtitleFormat.SRT,
            '.vtt': SubtitleFormat.VTT,
            '.json': SubtitleFormat.JSON,
            '.sbv': SubtitleFormat.SBV,
            '.ass': SubtitleFormat.ASS
        }
        
        return format_map.get(ext, SubtitleFormat.SRT)
    
    def _parse_srt_time(self, time_str: str) -> float:
        """解析SRT时间格式"""
        pattern = r'(\d{2}):(\d{2}):(\d{2}),(\d{3})'
        match = re.match(pattern, time_str)
        
        if match:
            hours = int(match.group(1))
            minutes = int(match.group(2))
            seconds = int(match.group(3))
            milliseconds = int(match.group(4))
            
            return hours * 3600 + minutes * 60 + seconds + milliseconds / 1000.0
        
        return 0.0
    
    def _parse_vtt_time(self, time_str: str) -> float:
        """解析VTT时间格式"""
        pattern = r'(\d{2}):(\d{2}):(\d{2})\.(\d{3})'
        match = re.match(pattern, time_str)
        
        if match:
            hours = int(match.group(1))
            minutes = int(match.group(2))
            seconds = int(match.group(3))
            milliseconds = int(match.group(4))
            
            return hours * 3600 + minutes * 60 + seconds + milliseconds / 1000.0
        
        return 0.0
    
    def _parse_sbv_time(self, time_str: str) -> float:
        """解析SBV时间格式"""
        pattern = r'(\d{1,2}):(\d{2}):(\d{2})\.(\d{3})'
        match = re.match(pattern, time_str)
        
        if match:
            hours = int(match.group(1))
            minutes = int(match.group(2))
            seconds = int(match.group(3))
            milliseconds = int(match.group(4))
            
            return hours * 3600 + minutes * 60 + seconds + milliseconds / 1000.0
        
        return 0.0
    
    def _parse_ass_time(self, time_parts: tuple) -> float:
        """解析ASS时间格式"""
        hours = int(time_parts[0])
        minutes = int(time_parts[1])
        seconds = int(time_parts[2])
        centiseconds = int(time_parts[3])
        
        return hours * 3600 + minutes * 60 + seconds + centiseconds / 100.0
    
    def _format_srt_time(self, seconds: float) -> str:
        """格式化SRT时间"""
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        secs = int(seconds % 60)
        milliseconds = int((seconds % 1) * 1000)
        
        return f"{hours:02d}:{minutes:02d}:{secs:02d},{milliseconds:03d}"
    
    def _format_vtt_time(self, seconds: float) -> str:
        """格式化VTT时间"""
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        secs = int(seconds % 60)
        milliseconds = int((seconds % 1) * 1000)
        
        return f"{hours:02d}:{minutes:02d}:{secs:02d}.{milliseconds:03d}"
    
    def _format_sbv_time(self, seconds: float) -> str:
        """格式化SBV时间"""
        return self._format_vtt_time(seconds)
    
    def _format_ass_time(self, seconds: float) -> str:
        """格式化ASS时间"""
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        secs = int(seconds % 60)
        centiseconds = int((seconds % 1) * 100)
        
        return f"{hours}:{minutes:02d}:{secs:02d}.{centiseconds:02d}"
    
    def _split_text(self, text: str) -> List[str]:
        """分割文本"""
        # 按句子分割
        sentences = re.split(r'[.!?。！？]+', text)
        sentences = [s.strip() for s in sentences if s.strip()]
        
        # 进一步分割长句
        result = []
        for sentence in sentences:
            if len(sentence) > 100:
                # 按标点或空格分割
                parts = re.split(r'[,;，；]\s*', sentence)
                for part in parts:
                    if part.strip():
                        result.append(part.strip())
            else:
                result.append(sentence)
        
        return result
    
    def _append_timelines(self, timelines: List[Timeline]) -> Timeline:
        """追加时间轴"""
        entries = []
        current_time = 0.0
        
        for timeline in timelines:
            for entry in timeline.entries:
                new_entry = SubtitleEntry(
                    index=len(entries) + 1,
                    start_time=current_time + entry.start_time,
                    end_time=current_time + entry.end_time,
                    text=entry.text,
                    metadata=entry.metadata
                )
                entries.append(new_entry)
            
            current_time += timeline.duration
        
        return Timeline(
            entries=entries,
            duration=current_time,
            metadata={'merged_mode': 'append'}
        )
    
    def _merge_timelines(self, timelines: List[Timeline]) -> Timeline:
        """合并时间轴"""
        all_entries = []
        
        for timeline in timelines:
            all_entries.extend(timeline.entries)
        
        # 按开始时间排序
        all_entries.sort(key=lambda x: x.start_time)
        
        return Timeline(
            entries=all_entries,
            duration=max((entry.end_time for entry in all_entries), default=0.0),
            metadata={'merged_mode': 'merge'}
        )
    
    def _overlay_timelines(self, timelines: List[Timeline]) -> Timeline:
        """叠加时间轴"""
        if not timelines:
            return Timeline(entries=[], duration=0.0)
        
        # 使用第一个时间轴作为基础
        base_timeline = timelines[0]
        
        return Timeline(
            entries=base_timeline.entries,
            duration=base_timeline.duration,
            metadata={'merged_mode': 'overlay'}
        )
